#include <chrono>
#include "asycipc/api.hpp"

#include <algorithm>
#include <limits>
#include <stdexcept>
#include <utility>
#include <unordered_map>

#include "asycipc/backpressure.hpp"
#include "asycipc/frame.hpp"
#include "asycipc/heartbeat.hpp"
#include "asycipc/sliding_window.hpp"
#include "asycipc/telemetry.hpp"

#include "asycipc/frame.hpp"

namespace asycipc {

namespace {

constexpr char kFragmentMagic[4] = {'F', 'R', 'A', 'G'};
constexpr uint8_t kFragmentHeaderVersion = 1;
constexpr std::size_t kFragmentHeaderSize = 4 + 1 + 1 + 2 + 2 + 8 + 2;  // magic + ver + reserved + idx + total + id + ext len

void AppendUint16(std::vector<uint8_t>& buffer, uint16_t value) {
  buffer.push_back(static_cast<uint8_t>(value & 0xFF));
  buffer.push_back(static_cast<uint8_t>((value >> 8) & 0xFF));
}

void AppendUint64(std::vector<uint8_t>& buffer, uint64_t value) {
  for (int i = 0; i < 8; ++i) {
    buffer.push_back(static_cast<uint8_t>((value >> (8 * i)) & 0xFF));
  }
}

uint16_t ReadUint16(const std::vector<uint8_t>& buffer, std::size_t offset) {
  return static_cast<uint16_t>(buffer[offset]) |
         static_cast<uint16_t>(buffer[offset + 1]) << 8;
}

uint64_t ReadUint64(const std::vector<uint8_t>& buffer, std::size_t offset) {
  uint64_t value = 0;
  for (int i = 0; i < 8; ++i) {
    value |= static_cast<uint64_t>(buffer[offset + i]) << (8 * i);
  }
  return value;
}

struct FragmentMetadata {
  uint64_t message_seq = 0;
  uint16_t total_chunks = 0;
  uint16_t chunk_index = 0;
  std::vector<uint8_t> user_extensions;
};

std::vector<uint8_t> BuildFragmentExtension(uint64_t message_seq, uint16_t total_chunks,
                                            uint16_t chunk_index,
                                            const std::vector<uint8_t>& user_extensions) {
  uint16_t ext_length = static_cast<uint16_t>(
      std::min<std::size_t>(user_extensions.size(), std::numeric_limits<uint16_t>::max()));
  std::vector<uint8_t> buffer;
  buffer.reserve(kFragmentHeaderSize + ext_length);
  buffer.insert(buffer.end(), std::begin(kFragmentMagic), std::end(kFragmentMagic));
  buffer.push_back(kFragmentHeaderVersion);
  buffer.push_back(0);  // reserved
  AppendUint16(buffer, chunk_index);
  AppendUint16(buffer, total_chunks);
  AppendUint64(buffer, message_seq);
  AppendUint16(buffer, ext_length);
  buffer.insert(buffer.end(), user_extensions.begin(), user_extensions.begin() + ext_length);
  return buffer;
}

std::optional<FragmentMetadata> ParseFragmentExtension(const Frame& frame) {
  if (frame.extensions.size() < kFragmentHeaderSize) {
    return std::nullopt;
  }
  if (!std::equal(std::begin(kFragmentMagic), std::end(kFragmentMagic),
                  frame.extensions.begin())) {
    return std::nullopt;
  }
  if (frame.extensions[4] != kFragmentHeaderVersion) {
    return std::nullopt;
  }
  const std::size_t base_offset = 6;  // magic(4) + version(1) + reserved(1)
  const uint16_t chunk_index = ReadUint16(frame.extensions, base_offset);
  const uint16_t total_chunks = ReadUint16(frame.extensions, base_offset + 2);
  const uint64_t message_seq = ReadUint64(frame.extensions, base_offset + 4);
  const uint16_t user_ext_len = ReadUint16(frame.extensions, base_offset + 12);
  const std::size_t required_size = kFragmentHeaderSize + user_ext_len;
  if (frame.extensions.size() < required_size) {
    return std::nullopt;
  }
  FragmentMetadata metadata;
  metadata.message_seq = message_seq;
  metadata.total_chunks = total_chunks;
  metadata.chunk_index = chunk_index;
  if (user_ext_len > 0) {
    metadata.user_extensions.assign(frame.extensions.begin() + kFragmentHeaderSize,
                                    frame.extensions.begin() + kFragmentHeaderSize + user_ext_len);
  }
  return metadata;
}

}  // namespace

// ---------------- PipeServer::Builder ----------------

PipeServer::Builder& PipeServer::Builder::Endpoint(std::string endpoint) {
  options_.endpoint = std::move(endpoint);
  return *this;
}

PipeServer::Builder& PipeServer::Builder::AllowApp(std::string app_id) {
  options_.allowed_app_ids.insert(std::move(app_id));
  return *this;
}

PipeServer::Builder& PipeServer::Builder::OnMessage(MessageHandler handler) {
  message_handler_ = std::move(handler);
  return *this;
}

PipeServer::Builder& PipeServer::Builder::OnError(ErrorHandler handler) {
  error_handler_ = std::move(handler);
  return *this;
}

PipeServer::Builder& PipeServer::Builder::OnBackpressure(
    BackpressureHandler handler) {
  backpressure_handler_ = std::move(handler);
  return *this;
}

PipeServer::Builder& PipeServer::Builder::ReliabilityEnabled(bool enabled) {
  options_.reliability_enabled = enabled;
  return *this;
}

PipeServer PipeServer::Builder::Build() const {
  return PipeServer(options_, message_handler_, error_handler_,
                    backpressure_handler_);
}

// ---------------- PipeServer ----------------

PipeServer::PipeServer(ServerOptions options, MessageHandler msg_handler,
                       ErrorHandler err_handler,
                       BackpressureHandler backpressure_handler)
    : options_(std::move(options)),
      message_handler_(std::move(msg_handler)),
      error_handler_(std::move(err_handler)),
      backpressure_handler_(std::move(backpressure_handler)) {
  handshake_config_.endpoint_name = options_.endpoint;
  handshake_config_.allowed_app_ids = options_.allowed_app_ids;
  handshake_config_.protocol_version = kProtocolVersion;
}

void PipeServer::Start() {
  handshake_manager_ = std::make_shared<HandshakeManager>(handshake_config_);
}

HandshakeResponse PipeServer::Accept(const HandshakeRequest& request,
                                     SessionContext& session) const {
  if (!handshake_manager_) {
    throw HandshakeError("server not started");
  }
  return handshake_manager_->Negotiate(request, session);
}

std::unique_ptr<Channel> PipeServer::CreateChannel(
    SessionContext session) const {
  if (!handshake_manager_) {
    throw HandshakeError("server not started");
  }

  auto channel = std::make_unique<Channel>(std::move(session), handshake_manager_);
  channel->SetBackpressureHandler(backpressure_handler_);
  channel->SetMessageHandler(message_handler_);
  channel->SetErrorHandler(error_handler_);

  if (options_.reliability_enabled) {
    const std::size_t window_size = options_.window_size == 0 ? 1 : options_.window_size;
    auto sliding_window = std::make_shared<SlidingWindow>(
        window_size, options_.retry_base, options_.max_retry_attempts);
    auto heartbeat = std::make_shared<HeartbeatMonitor>(
        options_.heartbeat_interval, options_.heartbeat_allowed_miss);
    const std::size_t low_watermark =
        window_size > 1 ? window_size / 2 : window_size;
    auto backpressure =
        std::make_shared<BackpressureController>(window_size, low_watermark);

    channel->SetReliabilityModules(sliding_window, heartbeat, backpressure,
                                   const_cast<TelemetryCenter*>(&telemetry_));
  }
  return channel;
}
std::vector<Frame> PipeServer::ProcessInbound(Channel& channel, const std::vector<uint8_t>& raw_payload) {
  Frame frame;
  try {
    frame = Frame::Deserialize(raw_payload);
  } catch (const FrameValidationError& ex) {
    if (error_handler_) {
      error_handler_(ex.what());
    }
    return {};
  }

  auto result = channel.HandleInboundFrame(frame);
  std::vector<Frame> frames = std::move(result.control_frames);

  if (result.application_frame.has_value() && message_handler_) {
    try {
      auto responses = message_handler_(channel, *result.application_frame);
      if (!responses.empty()) {
        frames.insert(frames.end(), responses.begin(), responses.end());
      }
    } catch (const std::exception& ex) {
      if (error_handler_) {
        error_handler_(ex.what());
      }
    } catch (...) {
      if (error_handler_) {
        error_handler_("unknown exception in message handler");
      }
    }
  }
  return frames;
}

std::vector<Frame> PipeServer::CollectMaintenanceFrames(Channel& channel) {
  std::vector<Frame> frames;

  auto retries = channel.PollRetries();
  for (const auto& retry : retries) {
    frames.push_back(retry.frame);
  }

  if (channel.ShouldSendHeartbeat()) {
    frames.push_back(Frame::MakeHeartbeat());
    channel.MarkHeartbeatSent();
  }

  if (channel.HeartbeatExpired()) {
    if (error_handler_) {
      error_handler_("heartbeat expired");
    }
  }

  return frames;
}
void PipeServer::ReportError(const std::string& message) {
  if (error_handler_) {
    error_handler_(message);
  }
  telemetry_.IncrementCounter("ipc_error_total");
}

// ---------------- PipeClient::Builder ----------------

PipeClient::Builder& PipeClient::Builder::Endpoint(std::string endpoint) {
  options_.endpoint = std::move(endpoint);
  return *this;
}

PipeClient::Builder& PipeClient::Builder::AppId(std::string app_id) {
  options_.app_id = std::move(app_id);
  return *this;
}

PipeClient::Builder& PipeClient::Builder::Capabilities(
    std::vector<std::string> caps) {
  options_.capabilities = std::move(caps);
  return *this;
}

PipeClient PipeClient::Builder::Build() const { return PipeClient(options_); }

// ---------------- PipeClient ----------------

PipeClient::PipeClient(ClientOptions options) : options_(std::move(options)) {}

HandshakeRequest PipeClient::BuildHandshakeRequest() const {
  HandshakeRequest request;
  request.app_id = options_.app_id;
  request.capabilities = options_.capabilities;
  return request;
}

// ---------------- Channel ----------------

Channel::Channel(SessionContext session,
                 std::shared_ptr<HandshakeManager> manager)
    : session_(std::move(session)),
      handshake_manager_(std::move(manager)) {
  time_provider_ = [] { return std::chrono::steady_clock::now(); };
}

void Channel::SetMessageHandler(MessageHandler handler) {
  message_handler_ = std::move(handler);
}

void Channel::SetErrorHandler(ErrorHandler handler) {
  error_handler_ = std::move(handler);
}

void Channel::SetReliabilityModules(std::shared_ptr<SlidingWindow> window,
                                    std::shared_ptr<HeartbeatMonitor> heartbeat,
                                    std::shared_ptr<BackpressureController> backpressure,
                                    TelemetryCenter* telemetry) {
  sliding_window_ = std::move(window);
  heartbeat_monitor_ = std::move(heartbeat);
  backpressure_controller_ = std::move(backpressure);
  telemetry_ = telemetry;

  if (backpressure_controller_) {
    backpressure_controller_->SetHandler([this](bool active) {
      NotifyBackpressure(active);
    });
    backpressure_controller_->Update(sliding_window_ ? sliding_window_->PendingCount() : 0);
  }

  if (sliding_window_ && telemetry_) {
    telemetry_->SetGauge("ipc_send_window_pending",
                         static_cast<double>(sliding_window_->PendingCount()));
  }
}
void Channel::SetTimeProvider(
    std::function<std::chrono::steady_clock::time_point()> provider) {
  time_provider_ = std::move(provider);
}

Channel::InboundFrameResult Channel::HandleInboundFrame(const Frame& frame) {
  InboundFrameResult result;

  auto mark_activity = [this]() {
    if (heartbeat_monitor_) {
      heartbeat_monitor_->MarkReceived(Now());
    }
    session_.Touch();
  };

  switch (frame.type) {
    case FrameType::kAck: {
      mark_activity();
      OnAck(frame.seq_no);
      break;
    }
    case FrameType::kNack: {
      mark_activity();
      OnNack(frame.seq_no);
      break;
    }
    case FrameType::kHeartbeat: {
      mark_activity();
      result.control_frames.push_back(Frame::MakeHeartbeatAck());
      break;
    }
    case FrameType::kHeartbeatAck: {
      MarkHeartbeatReceived();
      session_.Touch();
      break;
    }
    case FrameType::kData: {
      if (heartbeat_monitor_) {
        heartbeat_monitor_->MarkReceived(Now());
      }
      try {
        session_.RegisterInbound(frame.seq_no);
      } catch (const SessionStateError&) {
        if (telemetry_) {
          telemetry_->IncrementCounter("ipc_seq_error_total");
        }
        session_.Touch();
        result.control_frames.push_back(Frame::MakeNack(frame.seq_no, "unexpected_seq"));
        break;
      }
      if (telemetry_) {
        telemetry_->IncrementCounter("ipc_recv_total");
      }
      if (frame.flags & kFrameFlagAckRequired) {
        result.control_frames.push_back(Frame::MakeAck(frame.seq_no));
      }

      Frame data_frame = frame;
      if ((data_frame.flags & kFrameFlagFragment) != 0) {
        auto metadata = ParseFragmentExtension(data_frame);
        if (!metadata || metadata->total_chunks == 0) {
          if (error_handler_) {
            error_handler_("invalid fragment metadata");
          }
          break;
        }
        if (telemetry_) {
          telemetry_->IncrementCounter("ipc_fragment_rx_total");
        }

        auto& state = inbound_fragments_[metadata->message_seq];
        if (state.total_chunks == 0) {
          state.total_chunks = metadata->total_chunks;
          state.parts.resize(state.total_chunks);
          state.received.resize(state.total_chunks, false);
          state.received_count = 0;
          state.base_flags = data_frame.flags;
          state.reserved = data_frame.reserved;
          state.type = data_frame.type;
          state.extensions = metadata->user_extensions;
        } else if (state.total_chunks != metadata->total_chunks) {
          inbound_fragments_.erase(metadata->message_seq);
          if (error_handler_) {
            error_handler_("fragment total mismatch");
          }
          break;
        }

        if (metadata->chunk_index >= state.total_chunks) {
          inbound_fragments_.erase(metadata->message_seq);
          if (error_handler_) {
            error_handler_("fragment index out of range");
          }
          break;
        }

        if (!state.received[metadata->chunk_index]) {
          state.parts[metadata->chunk_index] = data_frame.payload;
          state.received[metadata->chunk_index] = true;
          state.received_count++;
          if (metadata->chunk_index == 0 && !metadata->user_extensions.empty()) {
            state.extensions = metadata->user_extensions;
          }
        }

        if (state.received_count < state.total_chunks) {
          break;
        }

        std::vector<uint8_t> combined;
        combined.reserve(data_frame.payload.size() * state.total_chunks);
        for (const auto& part : state.parts) {
          combined.insert(combined.end(), part.begin(), part.end());
        }

        uint8_t aggregated_flags = state.base_flags & kFrameFlagAckRequired;
        Frame assembled = Frame::MakeData(metadata->message_seq, combined, state.extensions,
                                          aggregated_flags);
        assembled.reserved = state.reserved;
        assembled.type = state.type;
        assembled.flags &= static_cast<uint8_t>(~kFrameFlagFragment);

        data_frame = std::move(assembled);
        inbound_fragments_.erase(metadata->message_seq);
      }

      bool handled = false;
      if (message_handler_) {
        try {
          auto responses = message_handler_(*this, data_frame);
          if (!responses.empty()) {
            result.control_frames.insert(result.control_frames.end(), responses.begin(),
                                         responses.end());
          }
          handled = true;
        } catch (const std::exception& ex) {
          handled = false;
          if (error_handler_) {
            error_handler_(ex.what());
          }
        } catch (...) {
          handled = false;
          if (error_handler_) {
            error_handler_("unknown exception in message handler");
          }
        }
      }
      if (!handled) {
        result.application_frame = data_frame;
      }
      break;
    }
    default: {
      mark_activity();
      if (message_handler_) {
        try {
          auto responses = message_handler_(*this, frame);
          if (!responses.empty()) {
            result.control_frames.insert(result.control_frames.end(), responses.begin(),
                                         responses.end());
          }
        } catch (const std::exception& ex) {
          if (error_handler_) {
            error_handler_(ex.what());
          }
        } catch (...) {
          if (error_handler_) {
            error_handler_("unknown exception in message handler");
          }
        }
      } else {
        result.application_frame = frame;
      }
      break;
    }
  }

  return result;
}

std::vector<Frame> Channel::MakeDataFrames(const std::vector<uint8_t>& payload,
                                           const std::vector<uint8_t>& extensions) {
  const std::size_t max_payload = kMaxFramePayloadBytes;
  const std::size_t total_chunks =
      payload.size() <= max_payload ? 1
                                    : (payload.size() + max_payload - 1) / max_payload;

  std::vector<Frame> frames;
  frames.reserve(total_chunks);

  uint64_t first_seq = 0;
  for (std::size_t i = 0; i < total_chunks; ++i) {
    const std::size_t offset = i * max_payload;
    const std::size_t length =
        std::min<std::size_t>(max_payload, payload.size() - offset);
    std::vector<uint8_t> chunk(payload.begin() + offset,
                               payload.begin() + offset + length);

    const uint64_t seq = session_.NextSendSeq();
    if (i == 0) {
      first_seq = seq;
    }

    uint8_t flags = sliding_window_ ? kFrameFlagAckRequired : 0;
    std::vector<uint8_t> chunk_extensions = extensions;
    if (total_chunks > 1) {
      flags |= kFrameFlagFragment;
      chunk_extensions = BuildFragmentExtension(
          first_seq, static_cast<uint16_t>(total_chunks), static_cast<uint16_t>(i),
          (i == 0) ? extensions : std::vector<uint8_t>{});
    }

    Frame frame = Frame::MakeData(seq, chunk, chunk_extensions, flags);

    if (sliding_window_) {
      sliding_window_->Push(frame, Now());
      if (telemetry_) {
        telemetry_->SetGauge("ipc_send_window_pending",
                             static_cast<double>(sliding_window_->PendingCount()));
      }
      if (backpressure_controller_) {
        backpressure_controller_->Update(sliding_window_->PendingCount());
      } else if (sliding_window_->IsFull()) {
        NotifyBackpressure(true);
      }
    }

    if (telemetry_) {
      telemetry_->IncrementCounter("ipc_send_total");
      if (total_chunks > 1) {
        telemetry_->IncrementCounter("ipc_fragment_tx_total");
      }
    }

    frames.push_back(std::move(frame));
  }

  return frames;
}

void Channel::OnAck(uint64_t seq) {
  bool ack_consumed = false;
  if (sliding_window_) {
    if (sliding_window_->MarkAck(seq)) {
      ack_consumed = true;
      if (telemetry_) {
        telemetry_->IncrementCounter("ipc_ack_total");
        telemetry_->SetGauge("ipc_send_window_pending",
                             static_cast<double>(sliding_window_->PendingCount()));
      }
      if (backpressure_controller_) {
        backpressure_controller_->Update(sliding_window_->PendingCount());
      }
    }
  } else {
    ack_consumed = true;
  }

  if (ack_consumed) {
    if (!sliding_window_ && telemetry_) {
      telemetry_->IncrementCounter("ipc_ack_total");
    }
    try {
      session_.AckReceived(seq);
    } catch (const SessionStateError&) {
    }
  }
}

void Channel::OnNack(uint64_t seq) {
  if (!sliding_window_) {
    return;
  }
  if (sliding_window_->MarkNack(seq, Now())) {
    if (telemetry_) {
      telemetry_->IncrementCounter("ipc_nack_total");
    }
  }
}

std::vector<SlidingWindow::RetryItem> Channel::PollRetries() {
  if (!sliding_window_) {
    return {};
  }
  auto retries = sliding_window_->CollectRetries(Now());
  if (telemetry_ && !retries.empty()) {
    telemetry_->IncrementCounter("ipc_retry_total",
                                 static_cast<uint64_t>(retries.size()));
  }
  if (backpressure_controller_) {
    backpressure_controller_->Update(sliding_window_->PendingCount());
  }
  if (telemetry_) {
    telemetry_->SetGauge("ipc_send_window_pending",
                         static_cast<double>(sliding_window_->PendingCount()));
  }
  return retries;
}

bool Channel::ShouldSendHeartbeat() const {
  if (!heartbeat_monitor_) {
    return false;
  }
  return heartbeat_monitor_->ShouldSend(Now());
}

void Channel::MarkHeartbeatSent() {
  if (!heartbeat_monitor_) {
    return;
  }
  heartbeat_monitor_->MarkSent(Now());
  if (telemetry_) {
    telemetry_->IncrementCounter("ipc_heartbeat_sent");
  }
}

void Channel::MarkHeartbeatReceived() {
  if (!heartbeat_monitor_) {
    return;
  }
  heartbeat_monitor_->MarkReceived(Now());
  if (telemetry_) {
    telemetry_->IncrementCounter("ipc_heartbeat_recv");
  }
}

bool Channel::HeartbeatExpired() const {
  if (!heartbeat_monitor_) {
    return false;
  }
  return heartbeat_monitor_->IsExpired(Now());
}

void Channel::SetBackpressureHandler(BackpressureHandler handler) {
  backpressure_handler_ = std::move(handler);
}

void Channel::NotifyBackpressure(bool active) const {
  if (telemetry_) {
    telemetry_->SetGauge("ipc_backpressure_state", active ? 1.0 : 0.0);
  }
  if (backpressure_handler_) {
    backpressure_handler_(active);
  }
}

std::chrono::steady_clock::time_point Channel::Now() const {
  return time_provider_ ? time_provider_() : std::chrono::steady_clock::now();
}

}  // namespace asycipc








